# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# A script to run multinode training with submitit.
# --------------------------------------------------------

import argparse
import os
import uuid
from pathlib import Path

import main_finetune_test as classification
import submitit


def parse_args():
    classification_parser = classification.get_args_parser()
    parser = argparse.ArgumentParser("Submitit for MAE finetune", parents=[classification_parser])
    parser.add_argument("--ngpus", default=8, type=int, help="Number of gpus to request on each node")
    parser.add_argument("--nodes", default=2, type=int, help="Number of nodes to request")
    parser.add_argument("--timeout", default=4320, type=int, help="Duration of the job")
    parser.add_argument("--job_dir", default="", type=str, help="Job dir. Leave empty for automatic.")

    parser.add_argument("--partition", default="learnfair", type=str, help="Partition where to submit")
    parser.add_argument("--use_volta32", action='store_true', help="Request 32G V100 GPUs")
    parser.add_argument('--comment', default="", type=str, help="Comment to pass to scheduler")
    return parser.parse_args()


def get_shared_folder() -> Path:
    """
    获取一个共享文件夹的路径。如果在 /checkpoint 目录下存在以用户为名的文件夹，
    则返回该文件夹的路径；如果不存在，则抛出异常。
    """
    user = os.getenv("USER")  # 获取当前用户的用户名
    if Path("/checkpoint/").is_dir():  # 检查 /checkpoint 目录是否存在
        p = Path(f"/checkpoint/{user}/experiments")  # 生成用户对应的共享文件夹路径
        p.mkdir(exist_ok=True)  # 如果文件夹不存在，则创建它
        return p
    raise RuntimeError("No shared folder available")  # 如果 /checkpoint 目录不存在，抛出异常


def get_init_file():
    """
    创建一个初始化文件，确保初始化文件的父目录已创建且初始化文件本身不存在。
    初始化文件路径包含唯一的 UUID，以确保每次生成的文件名不同。
    """
    os.makedirs(str(get_shared_folder()), exist_ok=True)  # 创建共享文件夹（如果不存在）
    init_file = get_shared_folder() / f"{uuid.uuid4().hex}_init"  # 使用 UUID 生成唯一的初始化文件名
    if init_file.exists():
        os.remove(str(init_file))  # 如果文件已经存在，删除它
    return init_file  # 返回初始化文件的路径


class Trainer(object):
    """
    定义训练器类，用于管理模型训练过程。
    """
    def __init__(self, args):
        """
        初始化训练器。

        参数:
            args: 命令行参数或训练配置。
        """
        self.args = args  # 保存传入的参数

    def __call__(self):
        """
        通过调用训练器类的实例，开始训练过程。
        """
        import main_finetune_test as classification  # 导入 finetune 测试脚本

        self._setup_gpu_args()  # 设置与 GPU 相关的参数
        classification.main(self.args)  # 调用分类任务的 main 函数

    def checkpoint(self):
        """
        处理检查点，检查是否有现有的检查点文件，如果存在则继续训练。
        返回一个延迟提交任务，用于重新排队训练任务。
        """
        import os
        import submitit

        # 获取初始化文件路径，并作为 dist_url
        self.args.dist_url = get_init_file().as_uri()
        checkpoint_file = os.path.join(self.args.output_dir, "checkpoint.pth")  # 检查点文件路径
        if os.path.exists(checkpoint_file):
            self.args.resume = checkpoint_file  # 如果有检查点文件，则恢复训练

        print("Requeuing ", self.args)  # 打印重新排队的信息
        empty_trainer = type(self)(self.args)  # 创建一个新的 Trainer 实例
        return submitit.helpers.DelayedSubmission(empty_trainer)  # 返回一个延迟提交任务

    def _setup_gpu_args(self):
        """
        设置与 GPU 相关的参数，如 GPU 设备编号、分布式训练的参数等。
        """
        import submitit
        from pathlib import Path

        job_env = submitit.JobEnvironment()  # 获取当前提交的作业环境信息
        # 更新输出目录为包含作业 ID 的路径，确保每个作业都有独立的输出目录
        self.args.output_dir = Path(str(self.args.output_dir).replace("%j", str(job_env.job_id)))
        self.args.log_dir = self.args.output_dir  # 日志目录与输出目录相同
        self.args.gpu = job_env.local_rank  # 当前进程的 GPU ID
        self.args.rank = job_env.global_rank  # 当前进程的全局 rank
        self.args.world_size = job_env.num_tasks  # 当前集群中的总任务数
        print(f"Process group: {job_env.num_tasks} tasks, rank: {job_env.global_rank}")  # 打印作业信息


def main():
    """
    主函数，用于设置训练任务并提交给执行器（AutoExecutor）。
    """
    args = parse_args()  # 解析命令行参数
    if args.job_dir == "":
        args.job_dir = get_shared_folder() / "%j"  # 如果没有指定 job_dir，则使用共享文件夹并根据作业 ID 命名

    # 根据作业 ID 设置文件夹，方便追踪实验
    executor = submitit.AutoExecutor(folder=args.job_dir, slurm_max_num_timeout=30)  # 创建执行器实例

    num_gpus_per_node = args.ngpus  # 每个节点的 GPU 数量
    nodes = args.nodes  # 节点数量
    timeout_min = args.timeout  # 超时时间（分钟）

    partition = args.partition  # 分区
    kwargs = {}  # 额外的参数
    if args.use_volta32:
        kwargs['slurm_constraint'] = 'volta32gb'  # 使用特定硬件（如 Volta 32GB GPU）
    if args.comment:
        kwargs['slurm_comment'] = args.comment  # 提交作业时的评论

    # 更新执行器的参数
    executor.update_parameters(
        mem_gb=40 * num_gpus_per_node,  # 每个节点的内存（GB）
        gpus_per_node=num_gpus_per_node,  # 每个节点的 GPU 数量
        tasks_per_node=num_gpus_per_node,  # 每个节点的任务数量，通常与 GPU 数量相同
        cpus_per_task=10,  # 每个任务的 CPU 核心数
        nodes=nodes,  # 节点数量
        timeout_min=timeout_min,  # 超时时间（分钟）
        # 以下是与集群相关的参数
        slurm_partition=partition,  # SLURM 分区
        slurm_signal_delay_s=120,  # SLURM 信号延迟时间（秒）
        **kwargs  # 其他额外的参数
    )

    executor.update_parameters(name="mae")  # 设置作业名称

    args.dist_url = get_init_file().as_uri()  # 设置分布式训练的 URL
    args.output_dir = args.job_dir  # 设置输出目录

    trainer = Trainer(args)  # 创建 Trainer 实例
    job = executor.submit(trainer)  # 提交训练作业

    print(job.job_id)  # 打印提交的作业 ID


if __name__ == "__main__":
    main()  # 调用主函数

